#ifndef AVL_TREE_H_
#define AVL_TREE_H_

#include <vector>
#include <memory>
#include <time.h> 

using namespace std;
/*
 * @brief AVL树
 */
typedef short HEIGHT_TYPE;

template<class T>
class AVLTreeNode{
protected:
	T val;
	shared_ptr<AVLTreeNode<T>> children[2]; 
	weak_ptr<AVLTreeNode<T>> father; //父节点
	HEIGHT_TYPE height;
	template<typename> friend class AVLTree;

public:
	AVLTreeNode(T val){
		this->val = val;	
		children[0] = children[1] = nullptr;
		height = 1;
	}
	
	HEIGHT_TYPE getHeight(){
		return height;
	}
	
	HEIGHT_TYPE balanceFactor(){
		HEIGHT_TYPE factor = 0;
		if(children[0] != nullptr){
			factor += children[0]->height;
		}
		if(children[1] != nullptr){
			factor -= children[1]->height;
		}
		return factor;
	}

	//更新树的高度
	void updateHeight(){
		HEIGHT_TYPE h = 0;	
		if(children[0] != nullptr && children[0]->height > h){
				h = children[0]->height;
		}
		if(children[1] != nullptr && children[1]->height > h){
				h = children[1]->height;
		}
		height = h + 1;
	}

	void traversal(vector<AVLTreeNode<T>> &nodes){
		if(children[0] != nullptr){
			children[0]->traversal(nodes);
		}
		nodes.push_back(*this);
		if(children[1] != nullptr){
			children[1]->traversal(nodes);
		}
	}

	T getVal(){
		return val;
	}
};

template<class T>
class AVLTree{
protected:
	shared_ptr<AVLTreeNode<T>> root;

	//重新平衡树
	void insertFix(shared_ptr<AVLTreeNode<T>> node){
		shared_ptr<AVLTreeNode<T>> x = node;
		while(x != nullptr){
			HEIGHT_TYPE balanceFactor = x->balanceFactor();
			x->updateHeight();
			if(balanceFactor < -1){
				if(x->children[1] != nullptr){
					if(x->children[1]->balanceFactor() < 0){
						leftRotate(x);
					}else{
						rightRotate(x->children[1]);
						leftRotate(x);	
					}
				}
			}else if(balanceFactor > 1){
				if(x->children[0] != nullptr){
					if(x->children[0]->balanceFactor() > 0){
						rightRotate(x);
					}else{
						leftRotate(x->children[0]);
						rightRotate(x);
					}
				}
			}
			x = x->father.lock();
		}
	}

public:

	AVLTree(){
		root = nullptr;
	}

	void insert(T val){
		shared_ptr<AVLTreeNode<T>> node = make_shared<AVLTreeNode<T>>(val);
		shared_ptr<AVLTreeNode<T>> cur = root, p = nullptr;
		while(cur != nullptr){
			p = cur;
			if(val < cur->getVal()){
				cur = cur->children[0];
			}else{
				cur = cur->children[1];
			}
		}
		if(p == nullptr){
			root = node;
		}else{
			node->father = p;
			if(val < p->getVal()){
				p->children[0] = node;
			}else{
				p->children[1] = node;
			}
			insertFix(node);
		}
	}

	// index = 0 for left rotate , index = 1 for right rotate
	void rotate(shared_ptr<AVLTreeNode<T>> x, int index){
        int otherIndex = index ^ 1;
		shared_ptr<AVLTreeNode<T>> y = x->children[otherIndex];
		if(y == nullptr) return ;
		if(x->father.lock() != nullptr){
			if(x->father.lock()->children[index] == x){
				x->father.lock()->children[index] = y;
			}else{
				x->father.lock()->children[otherIndex] = y;
			}
		}else{
			root = y;
		}
		y->father = x->father;
		x->father = y;
		x->children[otherIndex] = y->children[index];
		if(y->children[index] != nullptr){
			y->children[index]->father = x;
		}
		y->children[index] = x;
		//调整树的高度
		x->updateHeight();
		y->updateHeight();
		if(y->father.lock() != nullptr){
			y->father.lock()->updateHeight();
		}
	}

	//左旋
	void leftRotate(shared_ptr<AVLTreeNode<T>> x){
        rotate(x, 0);
        /*
		shared_ptr<AVLTreeNode<T>> y = x->children[1];
		if(y == nullptr) return ;
		if(x->father.lock() != nullptr){
			if(x->father.lock()->children[0] == x){
				x->father.lock()->children[0] = y;
			}else{
				x->father.lock()->children[1] = y;
			}
		}else{
			root = y;
		}
		y->father = x->father;
		x->father = y;
		x->children[1] = y->children[0];
		if(y->children[0] != nullptr){
			y->children[0]->father = x;
		}
		y->children[0] = x;
		//调整树的高度
		x->updateHeight();
		y->updateHeight();
		if(y->father.lock() != nullptr){
			y->father.lock()->updateHeight();
		}
        */
	}

	//右旋
	void rightRotate(shared_ptr<AVLTreeNode<T>> x){
        rotate(x, 1);
        /*
		shared_ptr<AVLTreeNode<T>> y = x->children[0];
		if(y == nullptr) return ;
		if(x->father.lock() != nullptr){
			if(x->father.lock()->children[1] == x){
				x->father.lock()->children[1] = y;
			}else{
				x->father.lock()->children[0] = y;
			}
		}else{
			root = y;
		}
		y->father = x->father;
		x->father = y;
		x->children[0] = y->children[1];
		if(y->children[1] != nullptr){
			y->children[1]->father = x;
		}
		y->children[1] = x;
		//调整树的高度
		x->updateHeight();
		y->updateHeight();
		if(y->father.lock() != nullptr){
			y->father.lock()->updateHeight();
		}
        */
	}

	//遍历树
	vector<AVLTreeNode<T>> traversal(){
		vector<AVLTreeNode<T>> nodes;
		if(root != nullptr){
			root->traversal(nodes);
		}
		return nodes;
	}

	//随机生成一颗树
	void intRandomTree(int n){
		vector<int> nums(n);
		for(int i = 0; i < n; i++){
			nums[i] = i + 1;
		}
		srand((unsigned)time(nullptr));
		int pos;
		for(int i = n - 1; i >= 0; i--){
			pos = random() % (i + 1);
			insert(nums[pos]);
			nums[pos] = nums[i];
		}
	}

	//打印树
	void printTree(){
		cout << "{" << endl;
		if(root != nullptr){
			vector<shared_ptr<AVLTreeNode<T>>> que;			
			vector<int> 	depth;
			shared_ptr<AVLTreeNode<T>> head;
			que.push_back(root);
			depth.push_back(1);
			int l, d;
			l = 0;
			// BFS遍历
			while(l < que.size()){
				head = que[l];
				d = depth[l];
				if(head->children[0] != nullptr){
					que.push_back(head->children[0]);
					depth.push_back(d + 1);
				}
				if(head->children[1] != nullptr){
					que.push_back(head->children[1]);
					depth.push_back(d + 1);
				}
				l++;	
			}
			d = 1;
			for(int i = 0; i < que.size(); i++){
				if(i > 0){
					if(depth[i] != d){
						cout << endl;
					}else{
						cout << " ";
					}
					d = depth[i];
				}
				cout << "(" ;
				if(que[i]->father.lock() != nullptr){
					cout << que[i]->father.lock()->getVal();	
				}else{
					cout << "nullptr";	
				}
				cout << "," << que[i]->getVal() << ")";
			}
		}
		cout << endl << "}" << endl;
	}

	//按递增顺序打印
	void printSerial(){
		vector<AVLTreeNode<int>> nodes = traversal();
		cout << "[" ;
		for(int i = 0; i < nodes.size(); i++){
			if(i > 0){
				cout <<",";
			}
			cout << nodes[i].getVal();
		}
		cout << "]";
	}

	//批量测试
	static void batchTest(){
		//int batchNum = 10000;
		int batchNum = 1;
        int num = 10000;
		for(int i = 0; i < batchNum; i++){
			AVLTree<int> avlTree;
			avlTree.intRandomTree(num);
			vector<AVLTreeNode<int>> nodes = avlTree.traversal();
			bool flag = true;
			for(int j = 1; j < nodes.size(); j++){
				if(nodes[j-1].getVal() > nodes[j].getVal()){
					flag = false;
					break;
				}
			}
			if(flag){
				cout << "index=" << i << ",success" << endl;
			}else{
				cout << "index=" << i << ",error ";
				avlTree.printSerial();
				cout << endl;
			}
		}
	}

	shared_ptr<AVLTreeNode<T>> getRoot(){
		return root;
	}

};

#endif
